import numpy as np
from osgeo import gdal, osr
import matplotlib.pyplot as plt
from tqdm import tqdm
import time
import os
from pylandstats import Landscape
# 使用PyLandStats计算GeoTIFF的景观指数，并保存为GeoTIFF文件。但是程序运算特别慢
class LandscapeMetrics:
    def __init__(self, image_path, window_scale=20, weights=[1/3, 1/3, 1/3], use_padding=True, output_dir=None):
        self.image_path = image_path
        self.window_scale = window_scale
        self.a = weights[0]
        self.b = weights[1]
        self.c = weights[2]
        self.use_padding = use_padding
        self.output_dir = output_dir if output_dir else os.getcwd()
        self.data = None
        self.resolution = None
        self.geotransform = None
        self.projection = None
        self.window_size = None
        self.padded_data = None
        self.nodata = None
        self.load_data()

    def load_data(self):
        """使用GDAL加载GeoTIFF数据并预处理。"""
        dataset = gdal.Open(self.image_path)
        if dataset is None:
            raise ValueError("无法打开文件")

        band = dataset.GetRasterBand(1)
        self.data = band.ReadAsArray().astype(float)

        self.geotransform = dataset.GetGeoTransform()
        self.resolution = self.geotransform[1]
        self.projection = dataset.GetProjection()

        self.nodata = band.GetNoDataValue()
        if self.nodata is not None:
            self.data[self.data == self.nodata] = np.nan

        rows, cols = self.data.shape
        self.window_size = (rows // self.window_scale, cols // self.window_scale)

        if self.use_padding:
            pad_rows = self.window_size[0] // 2
            pad_cols = self.window_size[1] // 2
            self.padded_data = np.pad(self.data, ((pad_rows, pad_rows), (pad_cols, pad_cols)), mode='constant',
                                      constant_values=np.nan)
        else:
            self.padded_data = self.data

        dataset = None

    def normalize(self, array):
        """将数组归一化到[0, 1]范围。"""
        if self.nodata is not None:
            array = np.where(array == self.nodata, np.nan, array)

        array_min = np.nanmin(array)
        array_max = np.nanmax(array)

        if array_max == array_min:
            return np.zeros_like(array)

        normalized_array = (array - array_min) / (array_max - array_min)

        if self.nodata is not None:
            normalized_array = np.where(np.isnan(normalized_array), self.nodata, normalized_array)

        return normalized_array

    def calculate_landscape_stability(self, fi, si, fdi):
        """计算景观稳定性指数 (LS)。"""
        LS = 1 - (self.a * fi + self.b * si + self.c * fdi)
        return LS

    def calculate_window_metrics(self):
        """使用移动窗口法计算景观指数，并保存为GeoTIFF文件。"""
        rows, cols = self.data.shape
        window_rows, window_cols = self.window_size
        pad_rows = window_rows // 2 if self.use_padding else 0
        pad_cols = window_cols // 2 if self.use_padding else 0

        fi_array = np.zeros((rows, cols), dtype=np.float32)
        si_array = np.zeros_like(fi_array)
        fdi_array = np.zeros_like(fi_array)
        pd_array = np.zeros_like(fi_array)

        start_time = time.time()

        for i in tqdm(range(pad_rows, rows + pad_rows), desc="Processing rows"):
            for j in range(pad_cols, cols + pad_cols):
                window = self.padded_data[i - pad_rows:i + pad_rows + 1, j - pad_cols:j + pad_cols + 1]

                # 跳过全为 NaN 的窗口
                if np.all(np.isnan(window)):
                    fi, si, fdi, pd = np.nan, np.nan, np.nan, np.nan
                else:
                    # 确保数据是整数类型（分类数据）
                    window = window.astype(int)
                    landscape = Landscape(window, res=(self.resolution, self.resolution))
                    metrics = landscape.compute_landscape_metrics_df()
                    fi = metrics['patch_density'].values[0]  # 破碎度指数
                    si = metrics['edge_density'].values[0]  # 分离度指数
                    fdi = metrics['fractal_dimension_am'].values[0]  # 分维度指数
                    pd = metrics['patch_density'].values[0]  # 斑块密度指数

                fi_array[i - pad_rows, j - pad_cols] = fi
                si_array[i - pad_rows, j - pad_cols] = si
                fdi_array[i - pad_rows, j - pad_cols] = fdi
                pd_array[i - pad_rows, j - pad_cols] = pd

        end_time = time.time()
        print(f"计算完成，耗时: {end_time - start_time:.2f} 秒")

        fi_normalized = self.normalize(fi_array)
        si_normalized = self.normalize(si_array)
        fdi_normalized = self.normalize(fdi_array)
        pd_normalized = self.normalize(pd_array)

        ls_array = self.calculate_landscape_stability(fi_normalized, si_normalized, fdi_normalized)
        ls_normalized = self.normalize(ls_array)

        self.save_as_geotiff(fi_normalized, "fragmentation_index_normalized.tif")
        self.save_as_geotiff(si_normalized, "separation_index_normalized.tif")
        self.save_as_geotiff(fdi_normalized, "fractal_dimension_index_normalized.tif")
        self.save_as_geotiff(pd_normalized, "patch_density_index_normalized.tif")
        self.save_as_geotiff(ls_array, "landscape_stability_index.tif")
        self.save_as_geotiff(ls_normalized, "landscape_stability_index_normalized.tif")

        self.visualize_results(fi_normalized, "Fragmentation Index (Normalized)", "fragmentation_index_normalized.png")
        self.visualize_results(si_normalized, "Separation Index (Normalized)", "separation_index_normalized.png")
        self.visualize_results(fdi_normalized, "Fractal Dimension Index (Normalized)", "fractal_dimension_index_normalized.png")
        self.visualize_results(pd_normalized, "Patch Density Index (Normalized)", "patch_density_index_normalized.png")
        self.visualize_results(ls_array, "Landscape Stability Index", "landscape_stability_index.png")
        self.visualize_results(ls_normalized, "Landscape Stability Index (Normalized)", "landscape_stability_index_normalized.png")

    def save_as_geotiff(self, array, output_filename):
        """将数组保存为GeoTIFF文件。"""
        rows, cols = array.shape
        driver = gdal.GetDriverByName("GTiff")
        output_path = os.path.join(self.output_dir, output_filename)
        out_dataset = driver.Create(output_path, cols, rows, 1, gdal.GDT_Float32)
        out_dataset.SetGeoTransform(self.geotransform)
        out_dataset.SetProjection(self.projection)
        out_band = out_dataset.GetRasterBand(1)

        if self.nodata is not None:
            array[np.isnan(array)] = self.nodata
            out_band.SetNoDataValue(self.nodata)

        out_band.WriteArray(array)
        out_band.FlushCache()
        out_dataset = None

    def visualize_results(self, array, title, output_filename):
        """可视化景观指数结果并保存为PNG文件。"""
        output_path = os.path.join(self.output_dir, output_filename)
        plt.figure(figsize=(10, 10))
        plt.imshow(array, cmap='viridis', vmin=0, vmax=1)
        plt.colorbar(label=title)
        plt.title(title)
        plt.axis('off')
        plt.savefig(output_path, bbox_inches='tight', dpi=300)
        plt.close()


# 示例使用
if __name__ == "__main__":
    image_path = "D:/test/test/lu2022_1k.tif"
    output_path = "D:/test/test/LS"
    weights = [0.3, 0.4, 0.3]

    if not os.path.exists(output_path):
        os.makedirs(output_path)

    landscape = LandscapeMetrics(
        image_path,
        window_scale=20,
        weights=weights,
        use_padding=True,
        output_dir=output_path
    )

    landscape.calculate_window_metrics()